from mcp.server.fastmcp import FastMCP
from decimal import Decimal, getcontext

# 创建MCP服务实例
mcp_service = FastMCP("Math Service")

# 定义加法工具
@mcp_service.tool()
def add_numbers(a: float, b: float) -> str:
    """计算两个数的和"""
    return str(a + b)

# 定义减法工具
@mcp_service.tool()
def subtract_numbers(a: float, b: float) -> str:
    """计算两个数的差"""
    return str(a - b)

# 定义乘法工具
@mcp_service.tool()
def multiply_numbers(a: float, b: float) -> str:
    """计算两个数的积，使用高精度计算"""
    getcontext().prec = 50  # 设置高精度
    return str(Decimal(str(a)) * Decimal(str(b)))

# 定义除法工具
@mcp_service.tool()
def divide_numbers(a: float, b: float) -> str:
    """计算两个数的商"""
    try:
        if b == 0:
            return "错误：除数不能为零"
        return str(a / b)
    except Exception as e:
        return f"计算错误：{str(e)}"

# 定义查找第k个质数的工具
@mcp_service.tool()
def find_nth_prime(k: int) -> str:
    """查找第k个质数"""
    try:
        if k <= 0:
            return "错误：k必须是正整数"
        
        if k == 1:
            return "2"
        
        # 使用埃拉托斯特尼筛法
        limit = max(100, k * 20)  # 初始估计范围
        sieve = [True] * (limit + 1)
        sieve[0] = sieve[1] = False
        
        primes = []
        for num in range(2, limit + 1):
            if sieve[num]:
                primes.append(num)
                if len(primes) == k:
                    return str(num)
                for multiple in range(num * num, limit + 1, num):
                    sieve[multiple] = False
        
        # 如果范围不够，递归扩大范围
        return find_nth_prime(k)
    except Exception as e:
        return f"计算错误：{str(e)}"

if __name__ == "__main__":
    mcp_service.run(transport="stdio")